package com.mjk.common.web.getway.utils.loadbalancer.impl;

import com.mjk.common.web.getway.utils.loadbalancer.AbstractLoadBalancer;

import java.util.SortedMap;
import java.util.TreeMap;

/**
 * 一致性哈希算法
 */
public class ConsistentHashLoadBalance extends AbstractLoadBalancer {
    private final SortedMap<Integer, String> circle = new TreeMap<>();
    private final int replicas;

    public ConsistentHashLoadBalance(int replicas) {
        this.replicas = replicas;
    }

    public void addNode(String node) {
        for (int i = 0; i < replicas; i++) {
            int hash = getHash(node + i);
            circle.put(hash, node);
        }
    }

    public void removeNode(String node) {
        for (int i = 0; i < replicas; i++) {
            int hash = getHash(node + i);
            circle.remove(hash);
        }
    }

    public String getNextServer(String key) {
        if (circle.isEmpty()) {
            return null;
        }
        int hash = getHash(key);
        if (!circle.containsKey(hash)) {
            SortedMap<Integer, String> tailMap = circle.tailMap(hash);
            hash = tailMap.isEmpty() ? circle.firstKey() : tailMap.firstKey();
        }
        return circle.get(hash);
    }

    private int getHash(String key) {
        return key.hashCode();
    }
}
